from PyQt5.QtCore import QThread, pyqtSignal
import os
import numpy as np
from PIL import Image
from scipy.ndimage import median_filter

class BatchFlatfieldThread(QThread):
    """批量伪平场处理线程"""
    progress_signal = pyqtSignal(int, str)  # 发送当前处理的索引和文件名
    complete_signal = pyqtSignal()  # 处理完成信号
    error_signal = pyqtSignal(str)  # 错误信号
    
    def __init__(self, input_folder, output_folder, image_files):
        super().__init__()
        self.input_folder = input_folder
        self.output_folder = output_folder
        self.image_files = image_files
        self.running = True
    
    def run(self):
        try:
            from scipy.ndimage import median_filter
            import numpy as np
            from astropy.io import fits
            from PIL import Image
            
            # 确保输出目录存在
            os.makedirs(self.output_folder, exist_ok=True)
            
            # 处理每个文件
            for idx, filename in enumerate(self.image_files):
                if not self.running:
                    break
                
                # 发送进度信号
                self.progress_signal.emit(idx, filename)
                
                # 构建输入和输出路径
                input_path = os.path.join(self.input_folder, filename)
                output_filename = f"FW_{filename}"
                output_path = os.path.join(self.output_folder, output_filename)
                
                try:
                    # 根据文件类型处理
                    if filename.lower().endswith(('.fits', '.fit', '.fts')):
                        # 处理FITS文件
                        with fits.open(input_path) as hdul:
                            image = hdul[0].data
                            header = hdul[0].header
                            
                            # 处理图像
                            # 第1步：使用小窗口中值滤波识别噪点
                            small_median = median_filter(image, size=3)
                            
                            # 识别噪点
                            hot_pixels = image > small_median * 1.5
                            cold_pixels = image < small_median / 1.5
                            noise_pixels = hot_pixels | cold_pixels
                            
                            # 第2步：使用中值滤波估算背景趋势
                            background = median_filter(image, size=20)
                            
                            # 避免除零错误
                            background[background == 0] = 1e-10
                            
                            # 第3步：进行伪平场校正
                            processed_image = image / background
                            
                            # 第4步：对识别出的噪点应用中值滤波值来替换
                            processed_image[noise_pixels] = small_median[noise_pixels] / background[noise_pixels]
                            
                            # 保存处理后的图像
                            fits.writeto(output_path, processed_image, header, overwrite=True)
                            
                        # 继续处理下一个文件
                        continue
                    
                    # 处理非FITS图像文件
                    img = Image.open(input_path)
                    img_array = np.array(img)
                    
                    # 处理图像
                    processed_image = np.zeros_like(img_array, dtype=np.float32)
                    
                    # 根据图像类型进行处理
                    if len(img_array.shape) == 3:  # 彩色图像
                        # 分别处理每个颜色通道
                        for channel in range(img_array.shape[2]):
                            data = img_array[:,:,channel]
                            
                            # 第1步：使用小窗口中值滤波识别噪点
                            small_median = median_filter(data, size=3)
                            
                            # 识别噪点
                            hot_pixels = data > small_median * 1.5
                            cold_pixels = data < small_median / 1.5
                            noise_pixels = hot_pixels | cold_pixels
                            
                            # 第2步：使用中值滤波估算背景趋势
                            background = median_filter(data, size=20)
                            
                            # 避免除零错误
                            background[background == 0] = 1e-10
                            
                            # 第3步：进行伪平场校正
                            channel_processed = data / background
                            
                            # 第4步：对识别出的噪点应用中值滤波值来替换
                            channel_processed[noise_pixels] = small_median[noise_pixels] / background[noise_pixels]
                            
                            # 保存结果
                            processed_image[:,:,channel] = channel_processed
                    else:
                        # 处理普通图像文件
                        img = Image.open(input_path)
                        img_array = np.array(img).astype(float)
                        
                        if len(img_array.shape) == 3:
                            processed_image = np.zeros_like(img_array)
                            for channel in range(img_array.shape[2]):
                                data = img_array[:,:,channel]
                                small_median = median_filter(data, size=3)
                                hot_pixels = data > small_median * 1.5
                                cold_pixels = data < small_median / 1.5
                                noise_pixels = hot_pixels | cold_pixels
                                background = median_filter(data, size=20)
                                background[background == 0] = 1e-10
                                channel_processed = data / background
                                channel_processed[noise_pixels] = small_median[noise_pixels] / background[noise_pixels]
                                processed_image[:,:,channel] = channel_processed
                        else:
                            data = img_array
                            small_median = median_filter(data, size=3)
                            hot_pixels = data > small_median * 1.5
                            cold_pixels = data < small_median / 1.5
                            noise_pixels = hot_pixels | cold_pixels
                            background = median_filter(data, size=20)
                            background[background == 0] = 1e-10
                            processed_image = data / background
                            processed_image[noise_pixels] = small_median[noise_pixels] / background[noise_pixels]
                    
                    # 处理后的图像可能有值超出0-255范围，需要重新映射到0-255
                    min_val = np.min(processed_image)
                    max_val = np.max(processed_image)
                    
                    # 重新映射到0-255范围
                    if max_val > min_val:
                        processed_image = 255 * (processed_image - min_val) / (max_val - min_val)
                    
                    # 保存处理后的图像
                    result_img = Image.fromarray(processed_image.astype(np.uint8))
                    result_img.save(output_path)
                    
                except Exception as e:
                    self.error_signal.emit(f"处理图像 {filename} 时出错: {e}")
                    print(f"处理图像 {filename} 时出错: {e}，跳过此图像。")
                    continue
            
            # 发送完成信号
            if self.running:
                self.complete_signal.emit()
        
        except Exception as e:
            # 如果有未捕获的异常，发送错误信号
            self.error_signal.emit(str(e))
    
    def cancel(self):
        """取消处理操作"""
        self.running = False